import math
import numpy as np
import matplotlib.pyplot as plt

import gradient_decent
import new_ton
import BFGS

# 绘制等高线图
delta = 0.025
x = np.arange(-5.0, 5.0, delta)
y = np.arange(-5.0, 5.0, delta)
X, Y = np.meshgrid(x, y)
Z = (3 * X ** 2 + 2 * Y ** 2 + X * Y)

fig, ax = plt.subplots(figsize=(10, 10))

CS = ax.contour(X, Y, Z)
ax.clabel(CS, inline=True, fontsize=10)
ax.set_title('Gradient Descent\n3 * x ** 2 + 2 * y ** 2 + x * y', fontsize=15)


def f(X):
    return 3 * X[0] ** 2 + 2 * X[1] ** 2 + (X[2] - 1) ** 2
    # return 3 * X[0] ** 2 + 2 * X[1] ** 2 + X[0] * X[1]


def grad_f(X):
    return np.array([6 * X[0], 4 * X[1], 2 * (X[2] - 1)])
    # return np.array([6 * X[0] + X[1], 4 * X[1] + X[0]])


def hessian_f(X):
    # return np.mat([
    #     [6, 0, 0],
    #     [0, 4, 0],
    #     [0, 0, 2],
    # ])
    return np.mat([
        [6, 1],
        [1, 4],
    ])


X_cur = np.array([50.0, 5.0, 19])
# X_cur = np.array([5.0, 5.0])

# bfgs(X_cur)

print("初始点", X_cur)
# set_X = gradient_decent.gradient_decent(f, grad_f, X_cur)
# set_X = new_ton.new_ton(f, grad_f, hessian_f, X_cur)
set_X = BFGS.bfgs(f, grad_f, X_cur)

ax.plot(set_X[:, 0], set_X[:, 1], c='r', marker='o')
plt.show()
